package com.sx.gateway.filter;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.Refill;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author shengx
 * @date 2020/5/16 20:20
 */
public class IPRegisterGatewayFilter implements GatewayFilter, Ordered {
    private String test;
    // 令牌桶最大容量
    private int capacity;
    // 每次Token补充量
    private int refillTokens;
    // 补充时间间隔
    private int duration;

    private static final Map<String, Bucket> BUCKET_CACHE = new ConcurrentHashMap<>();
    private Bucket createNewBucket(){
        Refill refill = Refill.greedy(refillTokens, Duration.ofMinutes(duration));
        Bandwidth limit = Bandwidth.classic(capacity, refill);
        return Bucket4j.builder().addLimit(limit).build();
    }
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        String ip = exchange.getRequest().getRemoteAddress().getAddress().getHostAddress();
        Bucket bucket = BUCKET_CACHE.computeIfAbsent(ip, k-> createNewBucket());
//        System.out.println(test);
        System.out.println("IP:" + ip + ", has tokens:" + bucket.getAvailableTokens());
        if(bucket.tryConsume(1)){
            return chain.filter(exchange);
        }else{

            return authErro(exchange.getResponse(), "您频繁进⾏注册，请求已被拒绝");
        }
    }

    @Override
    public int getOrder() {
        return 0;
    }

    private Mono<Void> authErro(ServerHttpResponse resp, String mess) {
        resp.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        resp.getHeaders().add("Content-Type", "application/json;charset=UTF-8");
        String returnStr = mess;
        DataBuffer buffer = resp.bufferFactory().wrap(returnStr.getBytes(StandardCharsets.UTF_8));
        return resp.writeWith(Flux.just(buffer));
    }

    public void setTest(String test) {
        this.test = test;
    }

    public void setCapacity(int capacity) {
        this.capacity = capacity;
    }

    public void setRefillTokens(int refillTokens) {
        this.refillTokens = refillTokens;
    }

    public void setDuration(int duration) {
        this.duration = duration;
    }
}
